#include <gwbase/base/logging.h>
#include <gwbase/base/util.h>
#include <gwbase/base/condvar.h>
#include <gwbase/base/thread.h>
#include <surbasic/surbasic.h>
#include <gwbase/base/timestamp.h>
#include <fstream>
#include <string>
#include <getopt.h>

using namespace dicom;
using namespace SDBasic;

static int64_t surdoc_key = 0;
double total_size = 0.0;
int faild_count = 0;
double failed_size = 0.0;
double total_working_time_ms = 0.1;
dicom::Timestamp start_time;
dicom::Timestamp stop_time;
int work_done_count = 0;
dicom::Mutex mu;
dicom::CondVar more_work;
std::deque<std::string> work_queue;
FILE* file_id_file = NULL;

struct SingleFileStat {
  double interval_ms;
  double size;
  double speed_kbs;
};

static SingleFileStat GetSingleFileStat(dicom::Timestamp& start,
                                        dicom::Timestamp& end, size_t size) {
  SingleFileStat stat;
  stat.size = static_cast<double>(size);
  stat.interval_ms = static_cast<double>(end.MicroSecondsSinceEpoch() -
                                         start.MicroSecondsSinceEpoch()) /
                     1000.0;

  stat.speed_kbs = static_cast<double>(size) / stat.interval_ms;

  total_size += stat.size;
  total_working_time_ms += stat.interval_ms;
  return stat;
}

static struct option long_options[] = {{"user", required_argument, 0, 'u'},
                                       {"password", required_argument, 0, 'p'},
                                       {"server", required_argument, 0, 's'},
                                       {"threads", required_argument, 0, 't'},
                                       {"read", no_argument, 0, 'r'},
                                       {"debug", no_argument, 0, 'd'},
                                       {0, 0, 0, 0}};

struct BenchArgs {
  std::string user;
  std::string password;
  std::string server;
  int threads;
  bool is_read;
  std::vector<std::string> files;
  bool debug;
  std::string file_id_file;

  BenchArgs()
      : user("surdoctestooJUvF20@cc.cc"),
        password("aA111111"),
        server("https://open.surdoc.com"),
        threads(1),
        is_read(false),
        debug(true),
        file_id_file("file-id") {}

  void PrintArgs() const {
    fprintf(stderr, "-------------------------------------------\n");
    fprintf(stderr, "BENCH: Server:  \t\t%s\n", server.c_str());
    fprintf(stderr, "BENCH: User:    \t\t%s\n", user.c_str());
    fprintf(stderr, "BENCH: Password:\t\t%s\n", password.c_str());
    fprintf(stderr, "BENCH: Threads: \t\t%d\n", threads);
    fprintf(stderr, "BENCH: Read:    \t\t%s\n", (is_read ? "true" : "false"));
    fprintf(stderr, "BENCH: NumOfFiles:\t\t%ld\n", files.size());
    fprintf(stderr, "-------------------------------------------\n");
    ::fflush(stdout);
  }

  bool TryParseOption(int argc, char** argv) {
    int c;
    while (true) {
      int option_index = 0;
      c = ::getopt_long(argc, argv, "u:p:s:t:rdf:", long_options,
                        &option_index);
      if (c == -1) break;

      switch (c) {
        case 'u':
          this->user = optarg;
          break;
        case 'p':
          this->password = optarg;
          break;
        case 's':
          this->server = optarg;
          break;
        case 't':
          this->threads = ::atoi(optarg);
          break;
        case 'r':
          this->is_read = true;
          break;
        case 'd':
          this->debug = true;
          break;
        case 'f':
          this->file_id_file = optarg;
          break;
        default:
          fprintf(stderr, "invalid argument: %c\n", c);
          return false;
      }
    }

    for (int i = optind; i < argc; ++i) files.push_back(argv[i]);

    return true;
  }
};

int64_t DoLogin(const BenchArgs& conf) {
  AccessToken token;
  int64_t surdoc_key =
      SDBasic::Login(conf.server, conf.user, conf.password, token);
  if (surdoc_key <= 0) {
    fprintf(stderr, "BENCH: failed to login to surcloud\n");
    exit(-1);
  }

  return surdoc_key;
}

class BenchWorker {
 public:
  BenchWorker(int id, int64_t surdoc_key, bool is_read,
              const std::string& tmpdir)
      : id_(id),
        surdoc_key_(surdoc_key),
        is_read_(is_read),
        stop_(false),
        count_(1),
        output_dir_(tmpdir) {}

  ~BenchWorker() {}

  void Stop() { stop_ = true; }

  void Start() {
    dicom::ThreadPtr p(
        new dicom::Thread(boost::bind(&BenchWorker::DoWorkLoop, this),
                          StringPrintf("bench_worker.%02d", id_)));
    p->Start();
  }

 private:
  int DoWorkLoop() {
    while (!stop_) {
      ::fflush(stdout);
      std::string f;
      {
        LockGuard g(mu);
        while (!stop_ && work_queue.empty()) {
          more_work.Wait(&mu);
        }
        if (stop_) return 0;

        f = work_queue.front();
        work_queue.pop_front();
      }
      if (is_read_)
        TryDownloadFile(f);
      else
        TryUploadFile(f);
      work_done_count++;
      count_++;
      ::fflush(stdout);
    }
    return 0;
  }

  bool TryUploadFile(const std::string& file) {
    std::string path = StringPrintf("surcloud_bench-%05d.%02d.%04d.avi",
                                    ::getpid(), id_, count_);

    fprintf(stderr, "BENCH: Start uploading: %s -> %s\n", file.c_str(),
            path.c_str());
    dicom::Timestamp start = dicom::Timestamp::Now();
    tree::TreeEntity node;
    bool ok = SDBasic::Upload(surdoc_key_, file, "0", path,
                              SDBasic::time::Timestamp::Now(), node);

    if (!ok) {
      fprintf(stderr, "BENCH: Failed to upload: %s -> %s\n", file.c_str(),
              path.c_str());
      faild_count++;
      return false;
    }
    dicom::Timestamp end = dicom::Timestamp::Now();

    dicom::FileInfo info;
    Status s = dicom::StatPath(file, &info);
    if (!s.ok() || !info.is_regular()) {
      fprintf(stderr, "BENCH: Failed to upload (2): %s -> %s\n", file.c_str(),
              path.c_str());
      return false;
    }

    SingleFileStat stat = GetSingleFileStat(start, end, info.size);
    fprintf(
        stderr,
        "\nBENCH: Done upload: %s -> %s : %s @ %.2lf kB, %.2lf s, %.2lf kB/s\n",
        file.c_str(), path.c_str(), node.id.c_str(), stat.size / 1024.0,
        stat.interval_ms / 1000.0, stat.speed_kbs);

    fprintf(file_id_file, "%s\n", node.id.c_str());
    return true;
  }

  bool TryDownloadFile(const std::string& file_id) {
    std::string path = dicom::path::Join(
        output_dir_,
        StringPrintf("surcloud_bench-%05d.%02d.%04d", ::getpid(), id_, count_));

    fprintf(stderr, "BENCH: Start downloading: %s -> %s\n", file_id.c_str(),
            path.c_str());
    dicom::Timestamp start = dicom::Timestamp::Now();
    bool ok = SDBasic::Download(surdoc_key_, file_id, path);

    if (!ok) {
      fprintf(stderr, "BENCH: Failed to download: %s -> %s\n", file_id.c_str(),
              path.c_str());
      faild_count++;
      return false;
    }
    dicom::Timestamp end = dicom::Timestamp::Now();

    dicom::FileInfo info;
    Status s = dicom::StatPath(path, &info);
    if (!s.ok() || !info.is_regular()) {
      fprintf(stderr, "BENCH: Failed to download (2): %s -> %s\n",
              file_id.c_str(), path.c_str());
      return false;
    }

    SingleFileStat stat = GetSingleFileStat(start, end, info.size);

    fprintf(stderr, "BENCH: Done download: %s -> %s @ %.2lf kB, %.2lf kB/s\n",
            file_id.c_str(), path.c_str(), stat.size / 1024, stat.speed_kbs);
    return true;
  }

  int id_;
  int64_t surdoc_key_;
  bool is_read_;
  bool stop_;
  int count_;
  std::string output_dir_;
};

int main(int argc, char** argv) {
  BenchArgs conf;
  if (conf.TryParseOption(argc, argv) == false) {
    fprintf(stderr, "failed to parse arguments\n");
    return -1;
  }

  if (!conf.is_read) {
    file_id_file = fopen(conf.file_id_file.c_str(), "w");
    if (file_id_file == NULL) {
      fprintf(stderr, "BENCH: failed to open file id file: %s\n",
              conf.file_id_file.c_str());
      ::fflush(stdout);
      return -1;
    }
  } else {
    std::ifstream f(conf.file_id_file);
    while (!f.eof()) {
      std::string s;
      std::getline(f, s);
      if (s.empty()) continue;
      conf.files.push_back(s);
    }
    f.close();
  }

  if (conf.debug) conf.PrintArgs();

  surdoc_key = DoLogin(conf);
  fprintf(stderr, "BENCH: login to surcloud successfully\n");
  ::fflush(stdout);

  std::vector<BenchWorker*> workers;
  for (int i = 0; i < conf.threads; ++i) {
    BenchWorker* b = new BenchWorker(i + 1, surdoc_key, conf.is_read, "/tmp");
    workers.push_back(b);
    b->Start();
  }

  for (size_t i = 0; i < conf.files.size(); ++i) {
    work_queue.push_back(conf.files[i]);
    more_work.Broadcast();
  }

  while (work_done_count < static_cast<int>(conf.files.size())) {
    sleep(1);
  }

  fprintf(stderr, "-------------------------------------------\n");
  fprintf(stderr, "BENCH: %d Success, %d Failure\n",
          work_done_count - faild_count, faild_count);
  fprintf(stderr, "BENCH: Total size: %.2lfkB @ %.2lf s\n", total_size / 1024.0,
          total_working_time_ms / 1000.0);
  fprintf(stderr, "BENCH: Speed:      %.2lfkB/s (st), %.2lfkB/s (mt)\n",
          total_size / total_working_time_ms,
          total_size / total_working_time_ms * conf.threads);
  fprintf(stderr, "-------------------------------------------\n");
  ::fflush(stdout);
  more_work.Broadcast();

  for (int i = 0; i < conf.threads; ++i) {
    workers[i]->Stop();
  }

  return 0;
}
